import random
import gym
import numpy as np
import torch
import rl_utils


device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
seedseed = 0
random.seed(seedseed)
np.random.seed(seedseed)
torch.manual_seed(seedseed)

'''parameters'''
num_episodes = 500
hidden_dim = 128
discount_factor = 0.98
para_GAE_lmbda = 0.95
critic_lr = 1e-2
kl_constraint = 0.0005
alpha = 0.5

'''env'''
env_name = 'CartPole-v1'
env = gym.make(env_name)
env.reset(seed=seedseed)
state_dim = env.observation_space.shape[0]
action_dim = env.action_space.n
print(f'state_dim = {state_dim}')
print(f'action_dim = {action_dim}')


alg_name = 'ActorCritic'

if alg_name == 'TRPO':
    from on_policy.alg_TRPO_Discrete import TRPO
    agent = TRPO(hidden_dim, env.observation_space, env.action_space, para_GAE_lmbda, kl_constraint, alpha, critic_lr, discount_factor, device)
elif alg_name == 'PPO':
    actor_lr = 1e-3
    epochs = 10
    para_PPO_clip = 0.2
    from on_policy.alg_PPO_Discrete import PPO
    agent = PPO(state_dim, hidden_dim, action_dim, actor_lr, critic_lr, epochs, para_GAE_lmbda, para_PPO_clip, discount_factor, device)
elif alg_name == 'ActorCritic':
    actor_lr = 1e-3
    critic_lr = 1e-2
    num_episodes = 1000
    hidden_dim = 128
    from on_policy.alg_ActorCritic import ActorCritic
    agent = ActorCritic(state_dim, hidden_dim, action_dim, actor_lr, critic_lr,discount_factor, device)


print('Training!!!!')
return_list = rl_utils.train_on_policy_agent(env, agent, num_episodes)
rl_utils.plot_results(return_list, env_name, alg_name, string_train_test = 'Training', moving_average_weight = 9)

print('Testing!!!!')
return_list_test = rl_utils.test_agent(env, agent, num_episodes = 50)
rl_utils.plot_results(return_list_test, env_name, alg_name, string_train_test = 'Testing', moving_average_weight = 3)
print('Rendering!!!!')
rl_utils.test_agent_render(env, agent)







# time_start = time.perf_counter()  # 记录开始时间

# time_end = time.perf_counter()  # 记录结束时间
# time_sum = time_end - time_start  # 计算的时间差为程序的执行时间，单位为秒/s
# print('time = %f' %time_sum)